import numpy as np


class KalmanFilter:
    """
    简化版卡尔曼滤波器类，用于预测和更新目标的中心位置
    
    维护单个目标的状态估计，包括位置和速度
    支持高速目标跟踪丢失时的特殊处理
    """
    
    def __init__(self, dt=1.0, process_noise=0.1, measurement_noise=0.1, response_speed=1.0):
        """
        初始化卡尔曼滤波器
        
        Args:
            dt: 时间步长
            process_noise: 过程噪声协方差
            measurement_noise: 测量噪声协方差
            response_speed: 跟踪响应速度系数
        """
        self.dt = dt
        self.response_speed = response_speed
        
        # 状态转移矩阵 (x, y, vx, vy)
        self.F = np.array([
            [1, 0, dt, 0],
            [0, 1, 0, dt],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ])
        
        # 观测矩阵
        self.H = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0]
        ])
        
        # 过程噪声协方差矩阵
        self.Q = (process_noise * (1.0 + response_speed * 0.2)) * np.eye(4)
        velocity_noise_factor = 1.0 + response_speed * 0.1
        self.Q[2, 2] *= velocity_noise_factor  # vx噪声增强
        self.Q[3, 3] *= velocity_noise_factor  # vy噪声增强
        
        # 测量噪声协方差矩阵
        measurement_noise_factor = measurement_noise / (1.0 + response_speed * 0.1)
        self.R = measurement_noise_factor * np.eye(2)
        
        # 初始状态协方差矩阵
        initial_covariance = 10.0 + response_speed * 2.0
        self.P = initial_covariance * np.eye(4)
        
        # 状态向量 [x, y, vx, vy]
        self.x = np.zeros(4)
        self.initialized = False
        
        # 高速运动和跟踪丢失处理参数
        self.lost_tracking_count = 0
        self.max_speed_persistence = 25
        self.velocity_damping_factor = 0.95
        self.acceleration_compensation = 0.5
        
    def _initialize_state(self, center):
        """初始化状态向量"""
        self.x = np.array([center[0], center[1], 0, 0])
        self.initialized = True
        self._prev_center = center
    
    def predict(self):
        """
        预测下一个状态
        
        Returns:
            预测的中心坐标 (x, y) - 整数类型
        """
        if not self.initialized:
            raise ValueError("Kalman filter not initialized. Call update() first.")
        
        # 获取当前速度和速度大小
        vx, vy = self.x[2], self.x[3]
        speed_magnitude = np.sqrt(vx**2 + vy**2)
        
        # 基础状态预测
        predicted_state = np.dot(self.F, self.x)
        
        # 高速目标跟踪丢失处理 - 优化：调整高速运动标准为v>500，减少加速度影响
        if self.lost_tracking_count > 0:
            if speed_magnitude > 500.0:  # 高速运动（优化：调整为v>500算高速运动）
                # 动态调整最大保持帧数 - 优化：降低速度影响系数
                adaptive_max_persistence = self.max_speed_persistence + int(min(25, speed_magnitude * 0.01))  # 减少速度对持久帧数的影响
                
                if self.lost_tracking_count <= adaptive_max_persistence:
                    # 动态衰减因子 - 优化：调整衰减速率
                    base_decay = 0.01 if speed_magnitude < 500.0 else 0.005  # 降低衰减速率使高速目标保持稳定
                    velocity_factor = 1.0 - (base_decay * self.lost_tracking_count)
                    predicted_state[2] = vx * velocity_factor
                    predicted_state[3] = vy * velocity_factor
                    
                    # 加速度补偿 - 优化：大幅减少加速度项的影响
                    if speed_magnitude > 500.0:  # 仅对高速运动应用
                        # 降低加速度补偿系数，减少加速度项影响
                        comp_factor = min(0.2, self.acceleration_compensation * (speed_magnitude / 1000.0))  # 降低至原来的1/5以下
                        predicted_state[2] += vx * comp_factor * 0.3  # 再降低3倍影响
                        predicted_state[3] += vy * comp_factor * 0.3
                else:
                    # 自适应阻尼
                    adaptive_damping = self.velocity_damping_factor + min(0.05, (speed_magnitude / 100.0))
                    predicted_state[2] = vx * adaptive_damping
                    predicted_state[3] = vy * adaptive_damping
            else:
                # 低速运动阻尼
                predicted_state[2] = vx * self.velocity_damping_factor
                predicted_state[3] = vy * self.velocity_damping_factor
            
            self.lost_tracking_count += 1
        
        # 大加速度处理 - 新增：添加a>500时的特殊处理，减少加速度影响
        # 计算加速度大小
        acceleration_magnitude = np.sqrt((predicted_state[2] - vx)**2 + (predicted_state[3] - vy)**2) / self.dt if self.dt > 0 else 0
        if acceleration_magnitude > 500.0:  # 大加速度情况（a>500）
            # 对于大加速度，减小加速度的影响，保持运动稳定性
            # 重置速度到更合理范围
            max_safe_speed = 500.0  # 设置安全速度上限
            if speed_magnitude > max_safe_speed:
                scale_factor = max_safe_speed / speed_magnitude
                predicted_state[2] = vx * scale_factor  # 使用原始速度进行缩放
                predicted_state[3] = vy * scale_factor
        
        # 更新状态
        self.x = predicted_state
        
        # 预测协方差更新 - 优化：调整高速运动判断标准为v>500
        if self.lost_tracking_count > 0 and speed_magnitude > 500.0:  # 统一使用v>500作为高速标准
            F_temp = self.F.copy()
            # 优化：降低速度对协方差的影响系数
            speed_factor = min(1.2, 1.0 + speed_magnitude * 0.001)  # 从0.02降至0.001，大幅减少影响
            F_temp[0, 2] *= speed_factor
            F_temp[1, 3] *= speed_factor
            self.P = np.dot(np.dot(F_temp, self.P), F_temp.T) + self.Q
        else:
            self.P = np.dot(np.dot(self.F, self.P), self.F.T) + self.Q
        
        return (int(round(self.x[0])), int(round(self.x[1])))
    
    def update(self, center):
        """
        使用测量值更新状态
        
        Args:
            center: 测量到的中心坐标 (x, y)
        
        Returns:
            更新后的中心坐标 (x, y) - 整数类型
        """
        # 重置跟踪丢失计数
        self.lost_tracking_count = 0
        
        # 初始化状态
        if not self.initialized:
            self._initialize_state(center)
            return (int(round(self.x[0])), int(round(self.x[1])))
        
        # 计算卡尔曼增益
        S = np.dot(np.dot(self.H, self.P), self.H.T) + self.R
        K = np.dot(np.dot(self.P, self.H.T), np.linalg.inv(S))
        
        # 状态更新
        y = np.array(center) - np.dot(self.H, self.x)
        updated_x = self.x + np.dot(K, y)
        
        # 增强响应速度
        if self.response_speed > 1.0:
            mix_factor = min(0.9, self.response_speed * 0.01)
            updated_x[0] = updated_x[0] * (1 - mix_factor) + center[0] * mix_factor
            updated_x[1] = updated_x[1] * (1 - mix_factor) + center[1] * mix_factor
            
            if hasattr(self, '_prev_center'):
                dx = center[0] - self._prev_center[0]
                dy = center[1] - self._prev_center[1]
                vel_update_factor = min(0.3, self.response_speed * 0.005)
                updated_x[2] = updated_x[2] * (1 - vel_update_factor) + (dx / self.dt) * vel_update_factor
                updated_x[3] = updated_x[3] * (1 - vel_update_factor) + (dy / self.dt) * vel_update_factor
        
        # 更新协方差
        I_KH = np.eye(4) - np.dot(K, self.H)
        self.P = np.dot(np.dot(I_KH, self.P), I_KH.T) + np.dot(np.dot(K, self.R), K.T)
        
        # 保存状态和历史位置
        self.x = updated_x
        self._prev_center = center
        
        return (int(round(self.x[0])), int(round(self.x[1])))
    
    @property
    def velocity(self):
        """获取当前速度向量"""
        return (self.x[2], self.x[3])

class KalmanFilterTracker:
    """
    卡尔曼滤波器跟踪器，为每个类别维护单独的卡尔曼滤波器
    
    只使用类别(cls)区分目标，不使用id
    提供update和predict方法，支持速度和加速度计算与平滑
    """
    
    def __init__(self, dt=1.0, process_noise=0.1, measurement_noise=0.1, response_speed=1.0, history=3):
        """
        初始化跟踪器
        
        Args:
            dt: 时间步长
            process_noise: 过程噪声协方差
            measurement_noise: 测量噪声协方差
            response_speed: 跟踪响应速度系数
            history: 平滑历史长度
        """
        # 按类别存储滤波器
        self.filters = {}
        
        # 历史数据存储
        self.prev_velocities = {}
        self.history_velocities = {}
        self.history_accelerations = {}
        self.prev_smoothed = {}
        
        # 参数设置
        self.history = max(1, history)
        self.dt = dt
        self.process_noise = process_noise
        self.measurement_noise = measurement_noise
        self.response_speed = response_speed
        
    def reset_tracking(self, cls):
        """
        重置特定类别的跟踪器
        
        Args:
            cls: 目标类别
        """
        if cls in self.filters:
            del self.filters[cls]
            del self.prev_velocities[cls]
            del self.history_velocities[cls]
            del self.history_accelerations[cls]
            # 清理平滑历史
            for key_suffix in ['_vx', '_vy', '_ax', '_ay']:
                key = f"{cls}{key_suffix}"
                if key in self.prev_smoothed:
                    del self.prev_smoothed[key]
    
    def get_active_classes(self):
        """获取所有活跃的跟踪类别"""
        return list(self.filters.keys())
    
    def _smooth_value(self, values, value_type, key):
        """
        使用指数移动平均平滑值序列
        
        Args:
            values: 历史值列表
            value_type: 值类型 ('velocity' 或 'acceleration')
            key: 用于存储上一次平滑结果的键
        
        Returns:
            平滑后的值
        """
        if not values:
            return 0
        
        # 变化阈值设置
        thresholds = {'velocity': 8, 'acceleration': 20}
        threshold = thresholds.get(value_type, 8)
        
        # 指数移动平均
        alpha = 0.7
        smoothed = values[0]
        for v in values[1:]:
            smoothed = alpha * v + (1 - alpha) * smoothed
        
        # 应用变化阈值
        if key in self.prev_smoothed and abs(smoothed - self.prev_smoothed[key]) < threshold:
            smoothed = self.prev_smoothed[key]
        
        self.prev_smoothed[key] = smoothed
        return smoothed
    
    def update(self, cls, center):
        """
        更新特定类别的目标中心位置
        
        Args:
            cls: 目标类别
            center: 测量到的中心坐标 (x, y)
        
        Returns:
            tuple: (position, velocity, acceleration, norm_velocity, norm_acceleration)
                position: 更新后的中心坐标 (x, y) - 整数类型
                velocity: 目标速度 (vx, vy) - 浮点数类型
                acceleration: 目标加速度 (ax, ay) - 浮点数类型
                norm_velocity: 目标速度的范数 - 整数类型
                norm_acceleration: 目标加速度的范数 - 整数类型
        """
        # 初始化滤波器和历史数据
        if cls not in self.filters:
            self.filters[cls] = KalmanFilter(
                dt=self.dt,
                process_noise=self.process_noise,
                measurement_noise=self.measurement_noise,
                response_speed=self.response_speed
            )
            self.prev_velocities[cls] = None
            self.history_velocities[cls] = {'x': [], 'y': []}
            self.history_accelerations[cls] = {'x': [], 'y': []}
        
        # 更新滤波器并获取位置
        position = self.filters[cls].update(center)
        
        # 获取当前速度并更新历史
        current_vx, current_vy = self.filters[cls].velocity
        self.history_velocities[cls]['x'].append(current_vx)
        self.history_velocities[cls]['y'].append(current_vy)
        
        # 限制历史长度
        if len(self.history_velocities[cls]['x']) > self.history:
            self.history_velocities[cls]['x'] = self.history_velocities[cls]['x'][-self.history:]
            self.history_velocities[cls]['y'] = self.history_velocities[cls]['y'][-self.history:]
        
        # 平滑速度
        smoothed_vx = self._smooth_value(self.history_velocities[cls]['x'], 'velocity', f"{cls}_vx")
        smoothed_vy = self._smooth_value(self.history_velocities[cls]['y'], 'velocity', f"{cls}_vy")
        velocity = (smoothed_vx, smoothed_vy)
        
        # 计算并平滑加速度
        if self.prev_velocities[cls] is not None:
            prev_vx, prev_vy = self.prev_velocities[cls]
            raw_ax = (current_vx - prev_vx) / self.dt
            raw_ay = (current_vy - prev_vy) / self.dt
            
            self.history_accelerations[cls]['x'].append(raw_ax)
            self.history_accelerations[cls]['y'].append(raw_ay)
            
            # 限制历史长度
            if len(self.history_accelerations[cls]['x']) > self.history:
                self.history_accelerations[cls]['x'] = self.history_accelerations[cls]['x'][-self.history:]
                self.history_accelerations[cls]['y'] = self.history_accelerations[cls]['y'][-self.history:]
            
            # 平滑加速度
            smoothed_ax = self._smooth_value(self.history_accelerations[cls]['x'], 'acceleration', f"{cls}_ax")
            smoothed_ay = self._smooth_value(self.history_accelerations[cls]['y'], 'acceleration', f"{cls}_ay")
            acceleration = (smoothed_ax, smoothed_ay)
        else:
            acceleration = (0.0, 0.0)
        
        # 保存当前速度作为历史
        self.prev_velocities[cls] = (current_vx, current_vy)
        
        # 计算范数
        norm_velocity = int((velocity[0]**2 + velocity[1]**2)**0.5)
        norm_acceleration = int((acceleration[0]**2 + acceleration[1]**2)**0.5)
        
        return position, velocity, acceleration, norm_velocity, norm_acceleration
    
    def predict(self, cls):
        """
        预测特定类别的目标中心位置
        
        Args:
            cls: 目标类别
        
        Returns:
            预测的中心坐标 (x, y) - 整数类型
        """
        if cls not in self.filters:
            raise ValueError(f"No filter found for class {cls}")
        
        # 增加跟踪丢失计数
        self.filters[cls].lost_tracking_count += 1
        
        # 预测位置
        return self.filters[cls].predict()
    
    def reset_tracking(self, cls):
        """
        重置特定类别的跟踪器
        
        Args:
            cls: 目标类别
        """
        if cls in self.filters:
            del self.filters[cls]
            del self.prev_velocities[cls]
            del self.history_velocities[cls]
            del self.history_accelerations[cls]
            # 清理平滑历史
            for key_suffix in ['_vx', '_vy', '_ax', '_ay']:
                key = f"{cls}{key_suffix}"
                if key in self.prev_smoothed:
                    del self.prev_smoothed[key]
    
    def get_active_classes(self):
        """获取所有活跃的跟踪类别"""
        return list(self.filters.keys())



